use std::collections::HashMap;
use std::sync::Arc;
use jsonwebtoken::{decode, decode_header, DecodingKey, Header, TokenData, Validation};
use jsonwebtoken::jwk::{Jwk, JwkSet};
use once_cell::sync::Lazy;
use tokio::sync::RwLock;
use url::Url;
use crate::oidc::{Claims, IdTokenClaims, OidcProviderMetadata};
use crate::support::error::{AppError, AppResult};

static JWKS_CACHE: Lazy<RwLock<HashMap<String, Jwk>>> = Lazy::new(|| RwLock::new(HashMap::new()));
#[derive(Debug,Clone)]
pub struct JwkService{
    pub issuer_url: Url,
    pub jwks_uri: Url,
    pub client_id: String,
}

impl JwkService {
    pub async fn new(metadata: &OidcProviderMetadata) -> AppResult<Arc<Self>> {
        let client = Arc::new(Self {
            issuer_url: Url::parse(&metadata.issuer_uri)?,
            jwks_uri:  Url::parse(&metadata.jwks_uri)?,
            client_id: crate::config::get().auth_config().client_id().to_string()
        });
        client.refresh_jwks().await?;
        Ok(client)
    }
    pub async fn get_jwk(&self,header: &Header) -> Option<Jwk> {
        let cache_read = JWKS_CACHE.read().await;
        let kid = header.kid.as_ref()?;
        cache_read.get(kid).cloned()
    }

    pub async fn refresh_jwks(&self) -> AppResult<()> {
        let jwks_set: JwkSet =reqwest::get(self.jwks_uri.clone())
            .await?
            .json()
            .await?;
        // 拿到写锁
        let mut guard_cache_write = JWKS_CACHE.write().await;
        guard_cache_write.clear();
        for jwk in jwks_set.keys {
            if let Some(kid) = jwk.common.key_id.clone() {
                guard_cache_write.insert(kid, jwk);
            }
        }
        tracing::info!("JWKS refreshed successfully. Cached {} keys.", guard_cache_write.len());
        Ok(())
    }

    pub async fn validate_token(&self,token: &str) -> AppResult<Claims> {
        //1. 从token拿出header
        let header = decode_header(token).map_err(|_| {
            AppError::AuthServerError("Invalid token header".to_string())
        })?;

        // 2. 从根据header的kid从Cache拿到jwk
        let jwk = match self.get_jwk(&header).await {
            Some(jwk) => jwk,
            None => {
                //不存在jwk，尝试刷新
                self.refresh_jwks().await.map_err(|_| {
                    AppError::AuthServerError("Refresh token error".to_string())
                })?;
                let Some(jwk_after_refresh) = self.get_jwk(&header).await else {
                    return Err(AppError::AuthServerError("Refresh token ineffective".to_string()));
                };
                jwk_after_refresh
            }
        };


        //3. 创建DecodingKey & Validation
        let decoding_key = DecodingKey::from_jwk(&jwk).map_err(|_| {
            AppError::AuthServerError("Failed to create decoding key from JWK".to_string())
        })?;
        tracing::info!("jwk {:#?}", jwk);


        let mut validation = Validation::new(header.alg);
        validation.validate_exp = true;
        validation.set_audience(&[self.client_id.clone()]);
        validation.set_issuer(&[self.issuer_url.clone()]);
        // validation.validate_nbf = true;
        // validation.set_required_spec_claims(&["exp", "iat", "iss", "aud", "sub", "typ"]);
        tracing::info!("validation : {:#?}",validation);
        //4. token验签
        let token_data = decode(token, &decoding_key, &validation).map_err(|_| {
            AppError::AuthServerError("Invalid token".to_string())
        })?;
        let claims = token_data.claims;
        Ok(claims)
    }

    pub async fn validate_id_token(&self,expected_nonce: &str,id_token: &str) -> AppResult<()> {
        let header = decode_header(id_token).map_err(|e| {
            tracing::error!("Failed to decode ID token header: {:?}", e);
            AppError::AuthServerError("Invalid ID token header".to_string())
        })?;
        let jwk = match self.get_jwk(&header).await {
            Some(jwk) => jwk,
            None => {
                self.refresh_jwks().await.map_err(|e| {
                    tracing::error!("Failed to refresh JWKS for ID token: {:?}", e);
                    AppError::AuthServerError("Refresh JWKS error for ID token".to_string())
                })?;
                self.get_jwk(&header).await.ok_or_else(|| {
                    AppError::AuthServerError("JWK not found after refresh for ID token".to_string())
                })?
            }
        };

        let decoding_key = DecodingKey::from_jwk(&jwk).map_err(|e| {
            tracing::error!("Failed to create decoding key from JWK for ID token: {:?}", e);
            AppError::AuthServerError("Failed to create decoding key for ID token".to_string())
        })?;

        let mut validation = Validation::new(header.alg);
        validation.validate_exp = true;
        validation.set_issuer(&[self.issuer_url.clone()]);
        validation.set_audience(&[self.client_id.clone()]);
        let id_token_data: TokenData<IdTokenClaims> = decode(id_token, &decoding_key, &validation)
            .map_err(|e| {
                tracing::error!("ID token validation failed (decode/std claims): {:?}", e);
                // Provide more specific error based on e.kind() if needed
                match e.kind() {
                    jsonwebtoken::errors::ErrorKind::InvalidToken => AppError::AuthServerError("ID Token: Invalid token format".to_string()),
                    jsonwebtoken::errors::ErrorKind::InvalidSignature => AppError::AuthServerError("ID Token: Invalid signature".to_string()),
                    jsonwebtoken::errors::ErrorKind::InvalidIssuer => AppError::AuthServerError("ID Token: Invalid issuer".to_string()),
                    jsonwebtoken::errors::ErrorKind::InvalidAudience => AppError::AuthServerError("ID Token: Invalid audience".to_string()),
                    jsonwebtoken::errors::ErrorKind::ExpiredSignature => AppError::AuthServerError("ID Token: Expired".to_string()),
                    jsonwebtoken::errors::ErrorKind::ImmatureSignature => AppError::AuthServerError("ID Token: Not yet valid (nbf)".to_string()),
                    _ => AppError::AuthServerError(format!("ID Token validation error: {}", e)),
                }
            })?;
        let claims = id_token_data.claims;
        match &claims.nonce {
            Some(nonce_in_token) => {
                if nonce_in_token != expected_nonce {
                    tracing::error!(
                        "Nonce mismatch. Expected: {}, Got: {}",
                        expected_nonce,
                        nonce_in_token
                    );
                    return Err(AppError::AuthFlowError("ID token nonce mismatch.".to_string()));
                }
            }
            None => {
                tracing::error!("Nonce missing in ID token, but was expected.");
                return Err(AppError::AuthFlowError("Nonce missing in ID token.".to_string()));
            }
        }
        Ok(())

    }
}